#include <Windows.h>
#include <iostream>


void RemoveASLR(char* filePath)
{
	DWORD fileAttributes = GetFileAttributesA(filePath);
	if (fileAttributes == INVALID_FILE_ATTRIBUTES)
		throw std::runtime_error("File doesnt exist.");

	HANDLE fileHandle = CreateFileA(filePath, GENERIC_READ, 0, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
	if (fileHandle == INVALID_HANDLE_VALUE)
		throw std::runtime_error("File handle could not be opened.");

	DWORD fileSize = GetFileSize(fileHandle, 0);
	if (fileSize == INVALID_FILE_SIZE || fileSize == 0)
		throw std::runtime_error("Unable to read file size.");

	char* fileBuffer = new char[fileSize];

	DWORD bytesRead = 0;
	BOOL fileRead = ReadFile(fileHandle, fileBuffer, fileSize, &bytesRead, 0);
	if (!fileRead || bytesRead != fileSize)
		throw std::runtime_error("Failed to read the file.");

	CloseHandle(fileHandle);

	IMAGE_DOS_HEADER* dosHeader = reinterpret_cast<IMAGE_DOS_HEADER*>(fileBuffer);
	if (dosHeader->e_magic != 0x5A4D)
		throw std::runtime_error("Not a valid PE file.");

	IMAGE_NT_HEADERS* ntHeaders = reinterpret_cast<IMAGE_NT_HEADERS*>(fileBuffer + dosHeader->e_lfanew);
	if (ntHeaders->OptionalHeader.Magic != 0x010B)
		throw std::runtime_error("Not a valid PE file.");

	ntHeaders->OptionalHeader.DllCharacteristics &= ~IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE;

	fileHandle = CreateFileA(filePath, GENERIC_WRITE, 0, 0, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, 0);
	if (fileHandle == INVALID_HANDLE_VALUE)
		throw std::runtime_error("File handle could not be opened.");

	DWORD bytesWritten = 0;
	BOOL fileWritten = WriteFile(fileHandle, fileBuffer, fileSize, &bytesWritten, 0);
	if (!fileWritten || bytesWritten != fileSize)
		throw std::runtime_error("Failed to write the file.");

	CloseHandle(fileHandle);
}



int main(int argc, char** argv)
{
	if (argc >= 2)
	{
		try
		{
			RemoveASLR(argv[1]);
			std::cout << "ASLR has been succesfuly disabled." << std::endl;
		}
		catch (std::runtime_error& error)
		{
			std::cout << "Error: " << error.what() << std::endl;
		}
	}
	else
	{
		std::cout << "Please specify the file path." << std::endl;
	}
}